# -*- coding: utf-8 -*-

import os
import re

from typing import Dict

from javalang.tree import MethodDeclaration, TryStatement

from conf.entity.module.grpc_config import GrpcConfig, GrpcMappingTypeEnum
from com.maycur.grpc.base.base_tools import append_import_to_line, enrich_nested_converter
from com.maycur.grpc.base.grpc_entity import GrpcFullMapping, GrpcMethod, GrpcMessage, GrpcService, \
    IceGrpcMethodGroup, IceGrpcMethodMap
from com.maycur.grpc.base.base_entity import JavaMethodBuilder, java_utils_dict
from com.maycur.grpc.base.java_loader import JavaLoader, JavaClassInfo
from com.maycur.grpc.convertor.exception_handler_generator import ExceptionHandlerGenerator
from com.maycur.grpc.convertor.java_grpc_converter import GrpcAndJavaConverter
from com.maycur.utils import JavaAstUtils
from com.maycur.utils.CommonUtils import *


class ServantTransfer(object):
    config: GrpcConfig
    mapping: GrpcFullMapping
    converter: GrpcAndJavaConverter
    ice_servant_loader: JavaLoader
    grpc_servant_loader: JavaLoader
    # used to access service_exception instance inside handler_generator
    grpc_exception_handler_generator: ExceptionHandlerGenerator
    grpc_exception_handler: JavaClassInfo
    # key: ice servant, value: grpc servant
    ice_grpc_mapping: Dict[str, str]
    ice_dict: Dict[str, JavaClassInfo]
    grpc_dict: Dict[str, JavaClassInfo]
    grpc_changed_dict: Dict[str, JavaClassInfo]
    result: bool

    def __init__(self, config: GrpcConfig, mapping: GrpcFullMapping, converter: GrpcAndJavaConverter):
        self.config = config
        self.mapping = mapping
        self.converter = converter
        self.ice_servant_loader = JavaLoader()
        self.grpc_servant_loader = JavaLoader()
        self.ice_grpc_mapping = {}
        self.ice_dict = {}
        self.grpc_dict = {}
        self.grpc_changed_dict = {}
        self.result = True

    def _load_ice_servant(self):
        project_path = self.config.ice_config.project_path
        ice_servant_path = f'{project_path}/{self.config.ice_config.servant_path}'
        ice_servant_path = re.sub('/+', '/', ice_servant_path)
        if not os.path.exists(ice_servant_path):
            print(f'ice_servant_path not exist, the path: {ice_servant_path}')
            self.result = False
            return
        for item in os.listdir(ice_servant_path):
            directory_item = f'{ice_servant_path}/{item}'
            if os.path.isdir(directory_item):
                self.ice_servant_loader.load_by_package(directory_item)
            else:
                self.ice_servant_loader.load_by_class(ice_servant_path, item)
        for item in self.ice_servant_loader.classes:
            if 'IceServant' in item.annotation_set:
                self.ice_dict[item.simple_name] = item

    def _load_grpc_servant(self):
        project_path = self.config.ice_config.project_path
        full_path = f'{self.config.proto_config.servant.root_path}/{self.config.proto_config.servant.servant_path}'
        grpc_servant_path = f'{project_path}/{full_path}'
        grpc_servant_path = re.sub('/+', '/', grpc_servant_path)
        if not os.path.exists(grpc_servant_path):
            print(f'grpc_servant_path not exist, the path: {grpc_servant_path}')
            self.result = False
            return
        for item in os.listdir(grpc_servant_path):
            directory_item = f'{grpc_servant_path}/{item}'
            if os.path.isdir(directory_item):
                self.grpc_servant_loader.load_by_package(item)
            else:
                self.grpc_servant_loader.load_by_class(grpc_servant_path, item)
        for item in self.grpc_servant_loader.classes:
            if 'GrpcService' in item.annotation_set:
                self.grpc_dict[item.simple_name] = item

    def _load_grpc_exception_handler(self):
        grpc_servant = self.config.proto_config.servant
        project_path = self.config.ice_config.project_path
        handler_path = project_path + '/' + grpc_servant.root_path + '/' + grpc_servant.exception_handler_path
        handler_path = re.sub('/+', '/', handler_path)
        handler_file = handler_path + '/' + grpc_servant.exception_handler_name + '.java'
        handler_file = re.sub('/+', '/', handler_file)
        if not os.path.exists(handler_path) or not os.path.exists(handler_file):
            self.grpc_exception_handler_generator = ExceptionHandlerGenerator(self.config)
            self.grpc_exception_handler = self.grpc_exception_handler_generator.generate()
        else:
            java_loader = JavaLoader()
            java_loader.load_by_class(handler_path, grpc_servant.exception_handler_name + '.java')
            self.grpc_exception_handler_generator = ExceptionHandlerGenerator(self.config, is_exist=True)
            self.grpc_exception_handler = java_loader.simple_name_dict[grpc_servant.exception_handler_name]

    def _map_ice_2_grpc(self):
        if not self.result:
            return
        # processing special servant mapping defined in config
        for key in self.config.mapping_config.grpc_2_ice:
            mapping_item = self.config.mapping_config.grpc_2_ice[key]
            if mapping_item.mapping_type == GrpcMappingTypeEnum.SERVANT:
                self.ice_grpc_mapping[mapping_item.ice] = mapping_item.grpc

        # processing default servant mapping
        for ice_key in self.ice_dict:
            if ice_key.endswith('Servant'):
                module = ice_key[:ice_key.rindex('Servant')]
                grpc_key = self._get_grpc_servant_by_module(module)
                if grpc_key in self.grpc_dict:
                    self.ice_grpc_mapping[ice_key] = grpc_key

    def _get_grpc_servant_by_module(self, module: str):
        return module + self.config.proto_config.servant.servant_suffix

    def _get_ice_servant_by_module(self, module: str):
        return module + self.config.ice_config.servant_suffix

    def _get_grpc_base_package(self):
        return f'com.maycur.grpc.{self.config.proto_config.package}'

    def _do_create_grpc_servant(self, grpc_service: GrpcService) -> JavaClassInfo:
        grpc_key = upper_first(grpc_service.module_name) + self.config.proto_config.service_suffix
        grpc_key += self.config.proto_config.message_prefix
        servant_config = self.config.proto_config.servant
        servant_path = self.config.proto_config.project_path
        servant_path += f'/{servant_config.root_path}/{servant_config.servant_path}/{grpc_key}.java'
        servant_package = f'{servant_config.root_package}.{servant_config.servant_path}'
        grpc_servant = JavaClassInfo(is_exist=False)
        grpc_servant.package = servant_package
        grpc_servant.path = servant_path
        grpc_servant.simple_name = grpc_key
        grpc_servant.full_name = f'{grpc_servant.package}.{grpc_servant.simple_name}'

        # the line under package
        import_append_idx = 1

        grpc_servant.lines.append(f'package {grpc_servant.package}; \n')
        # an empty line to protect class document
        grpc_servant.lines.append('\n')
        grpc_servant.lines.append(JavaAstUtils.make_document_comment('Maycur Tool') + '\n')
        grpc_servant.lines.append('@Slf4j\n')
        grpc_servant.lines.append('@GrpcService\n')
        append_import_to_line(grpc_servant, java_utils_dict['Slf4j'], import_append_idx)
        append_import_to_line(grpc_servant, java_utils_dict['GrpcService'], import_append_idx)
        grpc_servant.logger = None

        impl_full_name = f'com.maycur.grpc.{self.config.proto_config.package}.{grpc_key}.{grpc_service.name}ImplBase'
        servant_class_define = f'public class {grpc_key} extends {impl_full_name} ' + '{\n'
        grpc_servant.lines.append(servant_class_define)
        grpc_servant.lines.append('}')

        return grpc_servant

    def _do_create_grpc_servant_method(self, ice_servant: JavaClassInfo, grpc_servant: JavaClassInfo,
                                       method: GrpcMethod):
        grpc_base_package = f'com.maycur.grpc.{self.config.proto_config.package}'
        request_import = f'{grpc_base_package}.{method.request}'
        method_response = method.response
        if 'google.protobuf.Empty' in method.response:
            response_import = java_utils_dict['GrpcEmpty']
            method_response = 'Empty'
        else:
            response_import = f'{grpc_base_package}.{method.response}'
        method_param = f'{method.request} request, StreamObserver<{method_response}> responseObserver'

        # logger
        if grpc_servant.logger is not None:
            log_var = grpc_servant.logger.name
        else:
            log_var = 'log'

        # request & response message
        if method.request in self.mapping.grpc.struct_dict:
            grpc_request = self.mapping.grpc.struct_dict[method.request]
        if method.response in self.mapping.grpc.struct_dict:
            grpc_response = self.mapping.grpc.struct_dict[method.response]

        # the line under package
        import_append_idx = 1
        append_import_to_line(grpc_servant, request_import, import_append_idx)
        append_import_to_line(grpc_servant, response_import, import_append_idx)
        append_import_to_line(grpc_servant, java_utils_dict['StreamObserver'], import_append_idx)

        grpc_m = JavaMethodBuilder(4)
        grpc_m.add_line('@Override', 1)
        grpc_m.add_line(f'public void {method.name}({method_param})' + ' {', 1)

        # exception handler identity
        ex_handler_identity = f'"{method.name}"'
        if grpc_request is not None:
            for item in grpc_request.fields:
                if item.name == 'entCode':
                    ex_handler_identity = 'entCode'
                    grpc_m.add_line('String entCode = request.getEntCode();', 2)
                    break

        grpc_m.add_line(f'{self.grpc_exception_handler.simple_name} handler = () -> ' + '{', 2)
        append_import_to_line(grpc_servant, self.grpc_exception_handler.full_name, import_append_idx)

        # gRPC builder init
        if method_response != 'Empty':
            builder_init_line = f'{method.response}.Builder resBuilder = {method.response}.newBuilder();'
            grpc_m.add_line(builder_init_line, 3)

        # ice business codes
        if method.name in ice_servant.method_dict:
            ice_m = ice_servant.method_dict[method.name]
            code_lines = self._extract_codes_from_try_catch_block(ice_servant, ice_m)
            if len(code_lines) > 0:
                for line in code_lines:
                    grpc_m.add_line(self._get_optimized_ice_line(line, grpc_request, grpc_servant), 0)
            else:
                # skip empty method block
                if ice_m.l_line_idx - ice_m.f_line_idx >= 1:
                    for idx in range(ice_m.f_line_idx, ice_m.l_line_idx):
                        opt_line = ice_servant.lines[idx]
                        # skip empty line
                        if opt_line.strip() != '':
                            grpc_m.add_line(self._get_optimized_ice_line(opt_line, grpc_request, grpc_servant), 1)

        # gRPC builder post process
        if method_response != 'Empty' and grpc_response is not None:
            if len(grpc_response.fields) == 1:
                prev_line = grpc_m.lines[len(grpc_m.lines) - 1].line
                upper_field_name = upper_first(grpc_response.fields[0].name)
                if 'repeated' in grpc_response.fields[0].f_type:
                    append_import_to_line(grpc_servant, java_utils_dict['CollectionUtils'], import_append_idx)
                    grpc_m.add_line('if (CollectionUtils.isNotEmpty(res))' + ' {', 3)
                    grpc_m.add_line('res.forEach(item -> {', 4)
                    grpc_m.add_line(f'resBuilder.add{upper_field_name}(item);', 5)
                    grpc_m.add_line('});', 4)
                    grpc_m.add_line('}', 3)
                elif 'var res' in prev_line:
                    grpc_m.add_line(f'resBuilder.set{upper_field_name}(res);', 3)

        # an empty line
        grpc_m.add_line('', 0)

        # processing grpc servant response
        if method_response == 'Empty':
            grpc_res = 'null'
        else:
            grpc_res = 'resBuilder.build()'
        grpc_m.add_line(f'responseObserver.onNext({grpc_res});', 3)

        grpc_m.add_line('responseObserver.onCompleted();', 3)
        grpc_m.add_line('};', 2)
        grpc_m.add_line(f'handler.handleException({ex_handler_identity}, responseObserver, {log_var});', 2)
        grpc_m.add_line('}', 1)
        grpc_m.add_line('', 0)
        origin_class_close_line = grpc_servant.lines[grpc_servant.class_close_idx]
        grpc_servant.lines[grpc_servant.class_close_idx] = grpc_m.build() + origin_class_close_line

    @staticmethod
    def _extract_codes_from_try_catch_block(ice_servant: JavaClassInfo, m_declaration: MethodDeclaration) -> List[str]:
        res_codes = []
        if len(m_declaration.declaration.body) > 0:
            first_try_statement = m_declaration.declaration.body[0]
            if isinstance(first_try_statement, TryStatement):
                st_len = len(first_try_statement.block)
                if st_len > 0:
                    first_st = first_try_statement.block[0]
                    last_st = first_try_statement.block[st_len - 1]
                    for idx in range(first_st.position.line, last_st.position.line + 1):
                        res_codes.append(ice_servant.lines[idx - 1])
        return res_codes

    def _get_optimized_ice_line(self, line: str, grpc_request: GrpcMessage, grpc_servant: JavaClassInfo):
        # the line under package
        import_append_idx = 1

        # collect imports
        if 'Objects.' in line:
            append_import_to_line(grpc_servant, java_utils_dict['Objects'], import_append_idx)
        if 'StringUtils.' in line:
            append_import_to_line(grpc_servant, java_utils_dict['StringUtils'], import_append_idx)
        if 'List<' in line or 'List ' in line:
            append_import_to_line(grpc_servant, java_utils_dict['List'], import_append_idx)
        if 'ArrayList<' in line or 'ArrayList ' in line:
            append_import_to_line(grpc_servant, java_utils_dict['ArrayList'], import_append_idx)
        if 'Arrays.' in line:
            append_import_to_line(grpc_servant, java_utils_dict['Arrays'], import_append_idx)
        if 'Lists.' in line:
            append_import_to_line(grpc_servant, java_utils_dict['Lists'], import_append_idx)
        if 'Map<' in line or 'Map ' in line:
            append_import_to_line(grpc_servant, java_utils_dict['Map'], import_append_idx)
        if 'HashMap' in line:
            append_import_to_line(grpc_servant, java_utils_dict['HashMap'], import_append_idx)
        if 'Maps.' in line:
            append_import_to_line(grpc_servant, java_utils_dict['Maps'], import_append_idx)
        if 'Collectors.' in line:
            append_import_to_line(grpc_servant, java_utils_dict['Collectors'], import_append_idx)
        if 'Collections.' in line:
            append_import_to_line(grpc_servant, java_utils_dict['Collections'], import_append_idx)
        if 'Optional.' in line:
            append_import_to_line(grpc_servant, java_utils_dict['Optional'], import_append_idx)

        # find the first '(' and last ')', in order to process request parameter
        part_left = ''
        part_middle = line
        part_right = ''
        if '(' in line:
            bracket_start_idx = line.index('(')
            try:
                bracket_end_idx = line.rindex(')')
                part_left = line[:bracket_start_idx + 1]
                part_middle = line[bracket_start_idx + 1:bracket_end_idx]
                part_right = line[bracket_end_idx:]
            except ValueError:
                # '()' was split in two lines
                part_middle = line[bracket_start_idx + 1:]

        if len(grpc_request.fields) > 0:
            for item in grpc_request.fields:
                if item.name in part_middle:
                    getter = f'request.get{upper_first(item.name)}'
                    if 'repeated' in item.f_type:
                        getter += 'List'
                    getter += '()'
                    if f'{item.name}es' in part_middle:
                        part_middle = part_middle.replace(f'{item.name}es', getter)
                    elif f'{item.name}s' in part_middle:
                        part_middle = part_middle.replace(f'{item.name}s', getter)
                    else:
                        part_middle = part_middle.replace(item.name, getter)

        # combine again
        line = part_left + part_middle + part_right

        # convert toIce to gRPC newGrpcXxx(for example: newGrpcDepartment)
        if '.toIce(' in line:
            parts = line.split('.toIce')
            if len(parts) == 2:
                brackets = JavaAstUtils.find_bracket_idx(line)
                if len(brackets) == 2 and brackets[1] - brackets[0] > 1:
                    converter_module = JavaAstUtils.rightest_word(parts[0])
                    converter_method = f'newGrpc{converter_module}'
                    if converter_method in self.converter.method_converter_dict:
                        converter = self.converter.method_converter_dict[converter_method]
                        enrich_nested_converter(grpc_servant, converter)
                        append_import_to_line(grpc_servant, converter.full_name, import_append_idx)
                        the_var = line[brackets[0]+1:brackets[1]]

                        start_idx = line.index(converter_module)
                        end_idx = brackets[1]

                        grpc_m = lower_first(converter.simple_name) + f'.newGrpc{converter_module}({the_var})'
                        line = line[:start_idx] + grpc_m + line[end_idx + 1:]

        # convert ::toIce to gRPC newGrpcXxx
        elif '::toIce' in line:
            # lambda express
            end_idx = line.index('::toIce')
            for i in range(end_idx - 1, -1, -1):
                if line[i].isalpha() or line[i].isnumeric():
                    continue
                start_idx = i
                break
            converter_module = line[start_idx+1:end_idx]
            converter_method = f'newGrpc{converter_module}'
            if converter_method in self.converter.method_converter_dict:
                converter = self.converter.method_converter_dict[converter_method]
                enrich_nested_converter(grpc_servant, converter)
                append_import_to_line(grpc_servant, converter.full_name, import_append_idx)

                grpc_m = f'item -> {lower_first(converter.simple_name)}.newGrpc{converter_module}(item)'
                line = line[:start_idx + 1] + grpc_m + line[end_idx + len('::toIce'):]

        # convert new MIce to grpc newGrpcXxx
        elif 'new MIce' in line and 'Exception' not in line:
            start_idx = line.index('new MIce')
            end_idx = start_idx + len('new MIce')
            bracket_start_idx = end_idx
            for i in range(end_idx, len(line)):
                if line[i].isalpha() or line[i].isnumeric():
                    continue
                bracket_start_idx = i
                break
            bracket_end_idx = line.index(')', bracket_start_idx)
            the_var = line[bracket_start_idx+1:bracket_end_idx]
            converter_module = line[end_idx:bracket_start_idx]
            if converter_module.endswith('Dto'):
                converter_module = converter_module[:-3]
            if the_var != '' and converter_module != '':
                converter_module = upper_first(converter_module)
                converter_method = f'newGrpc{converter_module}'
                if converter_method in self.converter.method_converter_dict:
                    converter = self.converter.method_converter_dict[converter_method]
                    enrich_nested_converter(grpc_servant, converter)
                    append_import_to_line(grpc_servant, converter.full_name, import_append_idx)
                    grpc_m = f'{lower_first(converter.simple_name)}.{converter_method}({the_var})'

                    line = line[:start_idx] + grpc_m + line[bracket_end_idx + 1:]

        # convert <MIceXxx> to GrpcXxx
        elif '<MIce' in line:
            bracket_start_idx = line.index('<MIce')
            bracket_end_idx = line.index('>', bracket_start_idx)
            word = line[bracket_start_idx+1:bracket_end_idx]
            domain_name = word.replace('MIce', '')
            if domain_name.endswith('Dto'):
                domain_name = domain_name[:-3]
            grpc_name = f'Grpc{upper_first(domain_name)}'
            if grpc_name in self.mapping.grpc.struct_dict:
                grpc_full_name = f'{self._get_grpc_base_package()}.{grpc_name}'
                append_import_to_line(grpc_servant, grpc_full_name, import_append_idx)
                line = line[:bracket_start_idx+1] + grpc_name + line[bracket_end_idx:]

        # return expression
        if 'return ' in line and 'null' not in line:
            line = line.replace('return ', 'var res = ')
        return line

    def _do_transfer(self):
        if not self.result:
            return
        if len(self.mapping.ice.ice_2_grpc_services) > 0:
            for item in self.mapping.ice.ice_2_grpc_services:
                ice_key = self._get_ice_servant_by_module(item.module_name)
                if ice_key in self.ice_dict:
                    ice_servant = self.ice_dict[ice_key]
                    # existing gRPC servant
                    if ice_key in self.ice_grpc_mapping:
                        grpc_key = self.ice_grpc_mapping[ice_key]
                        if grpc_key in self.grpc_dict:
                            grpc_servant = self.grpc_dict[grpc_key]
                        else:
                            raise ValueError(f'gRPC servant not found: {grpc_key} ')
                    # new gRPC servant
                    else:
                        grpc_servant = self._do_create_grpc_servant(item)
                        grpc_key = grpc_servant.simple_name

                    method_group = IceGrpcMethodGroup(ice_key, grpc_servant.simple_name)
                    for method_item in item.fields:
                        if method_item.name not in grpc_servant.method_dict:
                            # need transfer grpc servant method
                            self._do_create_grpc_servant_method(ice_servant, grpc_servant, method_item)
                            self.grpc_changed_dict[grpc_key] = grpc_servant

                            # record ice and gRPC method mapping
                            method_map = IceGrpcMethodMap(method_item.name, method_item.name)
                            method_group.add_method(method_map)
                    if len(method_group.methods) > 0:
                        self.mapping.ice.ice_2_grpc_method_groups.append(method_group)
                else:
                    raise ValueError(f'ice servant not found: {ice_key}')

    def _do_flush_disk(self):
        if len(self.grpc_changed_dict) > 0:
            for key in self.grpc_changed_dict:
                grpc_servant = self.grpc_changed_dict[key]
                if grpc_servant.is_exist:
                    # update the grpc servant
                    with open(grpc_servant.path, 'w') as grpc_f:
                        grpc_f.writelines(grpc_servant.lines)
                else:
                    # create new
                    pre_check_path = grpc_servant.path[:grpc_servant.path.rindex('/')]
                    if not os.path.exists(pre_check_path):
                        os.makedirs(pre_check_path)
                    with open(grpc_servant.path, 'x') as create_f:
                        create_f.writelines(grpc_servant.lines)
        if not self.grpc_exception_handler.is_exist:
            # create new exception handler
            pre_check_path = self.grpc_exception_handler.path[:self.grpc_exception_handler.path.rindex('/')]
            if not os.path.exists(pre_check_path):
                os.makedirs(pre_check_path)
            with open(self.grpc_exception_handler.path, 'x') as handler_f:
                handler_f.writelines(self.grpc_exception_handler.lines)
        service_exception = self.grpc_exception_handler_generator.service_exception
        if not service_exception.is_exist:
            # create new service exception
            pre_check_path = service_exception.path[:service_exception.path.rindex('/')]
            if not os.path.exists(pre_check_path):
                os.makedirs(pre_check_path)
            with open(service_exception.path, 'x') as create_f:
                create_f.writelines(service_exception.lines)
        elif self.grpc_exception_handler_generator.service_exception_changed:
            # update existed service exception
            with open(service_exception.path, 'w') as update_f:
                update_f.writelines(service_exception.lines)

    def transfer(self):
        self._load_ice_servant()
        self._load_grpc_servant()
        self._load_grpc_exception_handler()
        self._map_ice_2_grpc()
        try:
            self._do_transfer()
            if not self.config.test_mode:
                self._do_flush_disk()
        except Exception as e:
            print('failed to transfer servant, reason: ', e)
